package com.leetcode.根据数据结构分类.堆;

import java.util.*;

/**
 * @author: xiaomi
 * @date: 2021/2/11
 * @description: 703. 数据流中的第 K 大元素
 * https://leetcode-cn.com/problems/kth-largest-element-in-a-stream/
 */
public class A_703_数据流中的第K大元素 {

    public static void main(String[] args) {
//test1();
//        test3();
//        test4();
        test5();
    }

    public static void test1() {
        // 1
        KthLargest kthLargest = new KthLargest(3, new int[]{4, 5, 8, 2});
        System.out.println(kthLargest.add(3));
        System.out.println(kthLargest.add(5));
        System.out.println(kthLargest.add(10));
        System.out.println(kthLargest.add(9));
        System.out.println(kthLargest.add(4));
    }

    public static void test2() {
        // 1
        KthLargest kthLargest = new KthLargest(1, new int[0]);
        System.out.println(kthLargest.add(-3));
        System.out.println(kthLargest.add(-2));
        System.out.println(kthLargest.add(-4));
        System.out.println(kthLargest.add(0));
        System.out.println(kthLargest.add(4));
    }

    public static void test3() {
        // 1
        KthLargest2 kthLargest = new KthLargest2(3, new int[]{4, 5, 8, 2});
        System.out.println(kthLargest.add(3));
        System.out.println(kthLargest.add(5));
        System.out.println(kthLargest.add(10));
        System.out.println(kthLargest.add(9));
        System.out.println(kthLargest.add(4));
    }

    public static void test4() {
        // 1
        KthLargest2 kthLargest = new KthLargest2(1, new int[0]);
        System.out.println(kthLargest.add(-3));
        System.out.println(kthLargest.add(-2));
        System.out.println(kthLargest.add(-4));
        System.out.println(kthLargest.add(0));
        System.out.println(kthLargest.add(4));
    }

    public static void test5() {
        // 1
        KthLargest2 kthLargest = new KthLargest2(2, new int[]{0});
        System.out.println(kthLargest.add(-1));
        System.out.println(kthLargest.add(1));
        System.out.println(kthLargest.add(-2));
        System.out.println(kthLargest.add(-4));
        System.out.println(kthLargest.add(3));
    }

    /**
     * 由于 可以用 Heap 实现 TOP K 的问题；
     * 那么可以用 TOP K range - TOP K-1  = ans
     * 但是这种只能适用于 nums.length > len 的情况
     * ---
     * 结果发现输入的参数并不理想：没有一开始 k 和 nums 就有意义。
     */
    static class KthLargest {
        /**
         * 容量为 K 的小顶堆
         */
        private PriorityQueue<Integer> smallTopKHeap;
        /**
         * 使用 set 进行维护，节省 heap 的遍历
         * 由于存在重复的元素，那么还是直接用 heap 的 iterator 了
         * --
         * 但是可以用 lessSet
         */
//        private HashSet<Integer> kSet = new HashSet<>();
        /**
         * 容量为 K-1 的小顶堆
         */
        private PriorityQueue<Integer> smallTopLessKHeap;
        /**
         *
         */
        private Set<Integer> lessSet = new HashSet<>();

        public KthLargest(int k, int[] nums) {
            int len = nums.length;
            if (len != 0) {
                smallTopLessKHeap = new PriorityQueue<>(k - 1);
            }
            smallTopKHeap = new PriorityQueue<>(k);
            if (len < k) {
                for (int i = 0; i < len; i++) {
                    smallTopLessKHeap.offer(nums[i]);
                    smallTopKHeap.offer(nums[i]);
                }
                Iterator<Integer> it = smallTopLessKHeap.iterator();
                while (it.hasNext()) {
                    Integer next = it.next();
                    lessSet.add(next);
                }
            }

            int l = k - 1;
            int i = 0;
            for (; i < l; i++) {
                smallTopLessKHeap.offer(nums[i]);
                smallTopKHeap.offer(nums[i]);
            }
            smallTopKHeap.offer(nums[i]);
            //对 smallTopLessKHeap 另外处理
            if (smallTopLessKHeap.peek() < nums[i]) {
                //如果堆顶元素小于 新的元素 -> 将堆顶元素弹出 + 进行替换
                smallTopLessKHeap.poll();
                smallTopLessKHeap.offer(nums[i]);
            }
            //此时已经完成了 heap 的初始化
            for (i++; i < len; i++) {
                if (smallTopLessKHeap.peek() < nums[i]) {
                    smallTopLessKHeap.poll();
                    smallTopLessKHeap.offer(nums[i]);
                }
                if (smallTopKHeap.peek() < nums[i]) {
                    smallTopKHeap.poll();
                    smallTopKHeap.offer(nums[i]);
                }
            }

            //此时已经实现了 topK  和 topK-1 的 小顶堆
            Iterator<Integer> it = smallTopLessKHeap.iterator();
            while (it.hasNext()) {
                Integer next = it.next();
                lessSet.add(next);
            }

        }

        /**
         * 跟 constructor 中类似进行维护
         *
         * @param val
         * @return
         */
        public int add(int val) {
            if (!smallTopLessKHeap.isEmpty() && smallTopLessKHeap.peek() < val) {
                Integer poll = smallTopLessKHeap.poll();
                lessSet.remove(poll);
                smallTopLessKHeap.offer(val);
                lessSet.add(val);
            }
            if (!smallTopKHeap.isEmpty() && smallTopKHeap.peek() < val) {
                smallTopKHeap.poll();
                smallTopKHeap.offer(val);
            }
            //找到不相同的元素
            int min = Integer.MAX_VALUE;
            Iterator<Integer> it = smallTopKHeap.iterator();
            while (it.hasNext()) {
                Integer v = it.next();
                if (!lessSet.contains(v)) {
                    return v;
                }
                min = Math.min(min, v);
            }
            return min;
        }
    }


    /**
     * 鉴于上述中构造参数并不合适，而且 TOP K 的 range 求解问题画蛇添足了！
     * --
     * 使用一个 Heap 然后每次去遍历 heap 即可。
     */
    static class KthLargest2 {
        /**
         * 容量为 K 的小顶堆
         */
        private PriorityQueue<Integer> smallTopKHeap;
        int count = 0;

        public KthLargest2(int k, int[] nums) {
            smallTopKHeap = new PriorityQueue<>(k);
            int len = nums.length;
            count = k;
            if (len < k) {
                //如果 nums 的长度 小于 k
                for (int i = 0; i < len; i++) {
                    smallTopKHeap.offer(nums[i]);
                }
                return;
            }
            int i = 0;
            for (; i < k; i++) {
                smallTopKHeap.offer(nums[i]);
            }
            for (; i < len; i++) {
                if (smallTopKHeap.peek() < nums[i]) {
                    smallTopKHeap.poll();
                    smallTopKHeap.offer(nums[i]);
                }
            }
        }

        public int add(int val) {
            if (smallTopKHeap.isEmpty() || smallTopKHeap.size()< count) {
                smallTopKHeap.offer(val);
            } else if (smallTopKHeap.size() == count && smallTopKHeap.peek() < val) {
                smallTopKHeap.poll();
                smallTopKHeap.offer(val);
            }
            return smallTopKHeap.peek();
        }
    }
}
